package com.smh.aiagent.config;

import lombok.Getter;
import lombok.Setter;
import org.postgresql.ds.PGSimpleDataSource;
import org.springframework.ai.embedding.EmbeddingModel;
import org.springframework.ai.vectorstore.pgvector.PgVectorStore;
import org.springframework.boot.context.properties.ConfigurationProperties;
import org.springframework.boot.context.properties.EnableConfigurationProperties;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.jdbc.core.JdbcTemplate;

import static org.springframework.ai.vectorstore.pgvector.PgVectorStore.PgDistanceType.COSINE_DISTANCE;
import static org.springframework.ai.vectorstore.pgvector.PgVectorStore.PgIndexType.HNSW;

/**
 * @author shiminghui
 * @date 2025/6/2 21:20
 * @description: 配置PG VectorStore(由于与mysql数据库冲突,所以就不使用SpringAi的自动配置了)
 */
@Configuration
@EnableConfigurationProperties(PgVectorStoreConfig.PgVectorStoreProperties.class)
public class PgVectorStoreConfig {


    @Bean
    public PgVectorStore pgVectorStore(EmbeddingModel dashscopeEmbeddingModel, PgVectorStoreProperties pgVectorStoreProperties) {

        PGSimpleDataSource pgSimpleDataSource = new PGSimpleDataSource();
        pgSimpleDataSource.setURL(pgVectorStoreProperties.url);
        pgSimpleDataSource.setUser(pgVectorStoreProperties.username);
        pgSimpleDataSource.setPassword(pgVectorStoreProperties.password);
        JdbcTemplate jdbcTemplate = new JdbcTemplate(pgSimpleDataSource);

        // 创建一个PG VectorStore实例
        return PgVectorStore
                .builder(jdbcTemplate, dashscopeEmbeddingModel)
                .dimensions(1536) // 设置嵌入向量维度
                .distanceType(COSINE_DISTANCE) // 设置距离类型
                .indexType(HNSW) // 设置索引类型
                .initializeSchema(true) // 设置是否初始化数据库
                .schemaName("public") // 设置框架名称
                .vectorTableName("vector_store") // 设置表名称
                .maxDocumentBatchSize(10000) // 设置批量插入的文档数量
                .build();
    }

    @Setter
    @Getter
    @ConfigurationProperties(prefix = "custom.datasource.pgsql")
    public static class PgVectorStoreProperties {
        private String url;
        private String username;
        private String password;
    }

}
